{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# HybridBayesTree"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<a href=\"https://colab.research.google.com/github/borglab/gtsam/blob/develop/gtsam/hybrid/doc/HybridBayesTree.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "tags": [
          "remove-cell"
        ]
      },
      "outputs": [],
      "source": [
        "try:\n",
        "    import google.colab\n",
        "    %pip install --quiet gtsam-develop\n",
        "except ImportError:\n",
        "    pass  # Not running on Colab, do nothing"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "A `HybridBayesTree` is the hybrid equivalent of a `gtsam.GaussianBayesTree`. It represents the result of **multifrontal** variable elimination on a `HybridGaussianFactorGraph`.\n",
        "\n",
        "Like a standard Bayes tree, it's a tree structure where each node is a 'clique'. However, in a `HybridBayesTree`, each clique contains a `gtsam.HybridConditional` representing $P(F_k | S_k)$, where $F_k$ are the frontal variables eliminated in that clique, and $S_k$ are the separator variables shared with the parent clique.\n",
        "\n",
        "In **hybrid** Bayes trees discrete variables can only occur as frontal variables if the separator is entirely discrete, i.e., we will never have the situation that a discrete variable is conditioned on a continuous variable.\n",
        "\n",
        "The joint probability distribution $P(X, M)$ encoded by the tree is the product of all clique conditionals:\n",
        "$$\n",
        "P(X, M) = \\prod_k P(F_k : S_k)\n",
        "$$\n",
        "Compared to a `HybridBayesNet` (from sequential elimination), a `HybridBayesTree` is (a) tree-structured, (b) often has a shallower structure. Both are advantageous for inference tasks like marginalization and incremental updates, especially in sparse problems common in robotics (SLAM)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [],
      "source": [
        "import gtsam\n",
        "import numpy as np\n",
        "\n",
        "from gtsam import (\n",
        "    HybridGaussianFactorGraph, \n",
        "    JacobianFactor, DecisionTreeFactor, HybridGaussianFactor,\n",
        "    DiscreteValues\n",
        ")\n",
        "from gtsam.symbol_shorthand import X, D\n",
        "\n",
        "import graphviz"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Creating a HybridBayesTree (via Elimination)\n",
        "\n",
        "`HybridBayesTree` objects are obtained by performing **multifrontal** elimination on a `HybridGaussianFactorGraph`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Original HybridGaussianFactorGraph:\n"
          ]
        },
        {
          "data": {
            "image/svg+xml": [
              "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
              "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
              " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
              "<!-- Generated by graphviz version 12.0.0 (0)\n",
              " -->\n",
              "<!-- Pages: 1 -->\n",
              "<svg width=\"206pt\" height=\"84pt\"\n",
              " viewBox=\"0.00 0.00 206.00 83.60\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
              "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 79.6)\">\n",
              "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-79.6 202,-79.6 202,4 -4,4\"/>\n",
              "<!-- var7205759403792793600 -->\n",
              "<g id=\"node1\" class=\"node\">\n",
              "<title>var7205759403792793600</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"27\" cy=\"-57.6\" rx=\"27\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"27\" y=\"-52.55\" font-family=\"Times,serif\" font-size=\"14.00\">d0</text>\n",
              "</g>\n",
              "<!-- factor0 -->\n",
              "<g id=\"node4\" class=\"node\">\n",
              "<title>factor0</title>\n",
              "<ellipse fill=\"black\" stroke=\"black\" cx=\"27\" cy=\"-1.8\" rx=\"1.8\" ry=\"1.8\"/>\n",
              "</g>\n",
              "<!-- var7205759403792793600&#45;&#45;factor0 -->\n",
              "<g id=\"edge1\" class=\"edge\">\n",
              "<title>var7205759403792793600&#45;&#45;factor0</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M27,-39.28C27,-25.77 27,-8.54 27,-3.96\"/>\n",
              "</g>\n",
              "<!-- factor2 -->\n",
              "<g id=\"node6\" class=\"node\">\n",
              "<title>factor2</title>\n",
              "<ellipse fill=\"black\" stroke=\"black\" cx=\"81\" cy=\"-1.8\" rx=\"1.8\" ry=\"1.8\"/>\n",
              "</g>\n",
              "<!-- var7205759403792793600&#45;&#45;factor2 -->\n",
              "<g id=\"edge4\" class=\"edge\">\n",
              "<title>var7205759403792793600&#45;&#45;factor2</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M41.47,-42.18C55.47,-28.24 75.23,-8.54 79.96,-3.84\"/>\n",
              "</g>\n",
              "<!-- var8646911284551352320 -->\n",
              "<g id=\"node2\" class=\"node\">\n",
              "<title>var8646911284551352320</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"171\" cy=\"-57.6\" rx=\"27\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"171\" y=\"-52.55\" font-family=\"Times,serif\" font-size=\"14.00\">x0</text>\n",
              "</g>\n",
              "<!-- factor1 -->\n",
              "<g id=\"node5\" class=\"node\">\n",
              "<title>factor1</title>\n",
              "<ellipse fill=\"black\" stroke=\"black\" cx=\"171\" cy=\"-1.8\" rx=\"1.8\" ry=\"1.8\"/>\n",
              "</g>\n",
              "<!-- var8646911284551352320&#45;&#45;factor1 -->\n",
              "<g id=\"edge2\" class=\"edge\">\n",
              "<title>var8646911284551352320&#45;&#45;factor1</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M171,-39.28C171,-25.77 171,-8.54 171,-3.96\"/>\n",
              "</g>\n",
              "<!-- factor3 -->\n",
              "<g id=\"node7\" class=\"node\">\n",
              "<title>factor3</title>\n",
              "<ellipse fill=\"black\" stroke=\"black\" cx=\"117\" cy=\"-1.8\" rx=\"1.8\" ry=\"1.8\"/>\n",
              "</g>\n",
              "<!-- var8646911284551352320&#45;&#45;factor3 -->\n",
              "<g id=\"edge5\" class=\"edge\">\n",
              "<title>var8646911284551352320&#45;&#45;factor3</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M156.53,-42.18C142.53,-28.24 122.77,-8.54 118.04,-3.84\"/>\n",
              "</g>\n",
              "<!-- var8646911284551352321 -->\n",
              "<g id=\"node3\" class=\"node\">\n",
              "<title>var8646911284551352321</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"99\" cy=\"-57.6\" rx=\"27\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"99\" y=\"-52.55\" font-family=\"Times,serif\" font-size=\"14.00\">x1</text>\n",
              "</g>\n",
              "<!-- var8646911284551352321&#45;&#45;factor2 -->\n",
              "<g id=\"edge3\" class=\"edge\">\n",
              "<title>var8646911284551352321&#45;&#45;factor2</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M93.3,-39.58C88.77,-26.03 82.94,-8.61 81.39,-3.97\"/>\n",
              "</g>\n",
              "<!-- var8646911284551352321&#45;&#45;factor3 -->\n",
              "<g id=\"edge6\" class=\"edge\">\n",
              "<title>var8646911284551352321&#45;&#45;factor3</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M104.7,-39.58C109.23,-26.03 115.06,-8.61 116.61,-3.97\"/>\n",
              "</g>\n",
              "</g>\n",
              "</svg>\n"
            ],
            "text/plain": [
              "<graphviz.sources.Source at 0x1075396a0>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Elimination Ordering: Position 0: x0, x1, d0\n",
            "\n",
            "\n",
            "Resulting HybridBayesTree:\n"
          ]
        },
        {
          "data": {
            "image/svg+xml": [
              "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
              "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
              " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
              "<!-- Generated by graphviz version 12.0.0 (0)\n",
              " -->\n",
              "<!-- Title: G Pages: 1 -->\n",
              "<svg width=\"82pt\" height=\"188pt\"\n",
              " viewBox=\"0.00 0.00 82.04 188.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
              "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 184)\">\n",
              "<title>G</title>\n",
              "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-184 78.04,-184 78.04,4 -4,4\"/>\n",
              "<!-- 7205759403792793600 -->\n",
              "<g id=\"node1\" class=\"node\">\n",
              "<title>7205759403792793600</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"37.02\" cy=\"-162\" rx=\"27\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"37.02\" y=\"-156.95\" font-family=\"Times,serif\" font-size=\"14.00\">d0</text>\n",
              "</g>\n",
              "<!-- 8646911284551352321 -->\n",
              "<g id=\"node2\" class=\"node\">\n",
              "<title>8646911284551352321</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"37.02\" cy=\"-90\" rx=\"37.02\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"37.02\" y=\"-84.95\" font-family=\"Times,serif\" font-size=\"14.00\">x1 : d0</text>\n",
              "</g>\n",
              "<!-- 7205759403792793600&#45;&gt;8646911284551352321 -->\n",
              "<g id=\"edge1\" class=\"edge\">\n",
              "<title>7205759403792793600&#45;&gt;8646911284551352321</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M37.02,-143.7C37.02,-136.41 37.02,-127.73 37.02,-119.54\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"40.52,-119.62 37.02,-109.62 33.52,-119.62 40.52,-119.62\"/>\n",
              "</g>\n",
              "<!-- 8646911284551352320 -->\n",
              "<g id=\"node3\" class=\"node\">\n",
              "<title>8646911284551352320</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"37.02\" cy=\"-18\" rx=\"37.02\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"37.02\" y=\"-12.95\" font-family=\"Times,serif\" font-size=\"14.00\">x0 : x1</text>\n",
              "</g>\n",
              "<!-- 8646911284551352321&#45;&gt;8646911284551352320 -->\n",
              "<g id=\"edge2\" class=\"edge\">\n",
              "<title>8646911284551352321&#45;&gt;8646911284551352320</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M37.02,-71.7C37.02,-64.41 37.02,-55.73 37.02,-47.54\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"40.52,-47.62 37.02,-37.62 33.52,-47.62 40.52,-47.62\"/>\n",
              "</g>\n",
              "</g>\n",
              "</svg>\n"
            ],
            "text/plain": [
              "<graphviz.sources.Source at 0x10605bd40>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# --- Create a HybridGaussianFactorGraph ---\n",
        "hgfg = HybridGaussianFactorGraph()\n",
        "dk0 = (D(0), 2) # Binary discrete variable\n",
        "\n",
        "# Prior on D0: P(D0=0)=0.6, P(D0=1)=0.4\n",
        "prior_d0 = DecisionTreeFactor([dk0], \"0.6 0.4\")\n",
        "hgfg.add(prior_d0)\n",
        "\n",
        "# Prior on X0: P(X0) = N(0, 1)\n",
        "prior_x0 = JacobianFactor(X(0), np.eye(1), np.zeros(1), gtsam.noiseModel.Isotropic.Sigma(1, 1.0))\n",
        "hgfg.add(prior_x0)\n",
        "\n",
        "# Conditional measurement on X1: P(X1 | D0)\n",
        "# Mode 0: P(X1 | D0=0) = N(1, 0.25)\n",
        "gf0 = JacobianFactor(X(1), np.eye(1), np.array([1.0]), gtsam.noiseModel.Isotropic.Sigma(1, 0.5))\n",
        "# Mode 1: P(X1 | D0=1) = N(5, 1.0)\n",
        "gf1 = JacobianFactor(X(1), np.eye(1), np.array([5.0]), gtsam.noiseModel.Isotropic.Sigma(1, 1.0))\n",
        "meas_x1_d0 = HybridGaussianFactor(dk0, [gf0, gf1])\n",
        "hgfg.add(meas_x1_d0)\n",
        "\n",
        "# Factor connecting X0 and X1: P(X1 | X0) = N(X0+1, 0.1)\n",
        "odom_x0_x1 = JacobianFactor(X(0), -np.eye(1), X(1), np.eye(1), np.array([1.0]), gtsam.noiseModel.Isotropic.Sigma(1, np.sqrt(0.1)))\n",
        "hgfg.add(odom_x0_x1)\n",
        "\n",
        "print(\"Original HybridGaussianFactorGraph:\")\n",
        "# hgfg.print()\n",
        "display(graphviz.Source(hgfg.dot()))\n",
        "\n",
        "# --- Elimination ---\n",
        "# Use default hybrid ordering (continuous first, then discrete)\n",
        "ordering = gtsam.HybridOrdering(hgfg)\n",
        "print(f\"\\nElimination Ordering: {ordering}\")\n",
        "\n",
        "# Perform multifrontal elimination\n",
        "hybrid_bayes_tree = hgfg.eliminateMultifrontal(ordering)\n",
        "\n",
        "print(\"\\nResulting HybridBayesTree:\")\n",
        "# hybrid_bayes_tree.print()\n",
        "display(graphviz.Source(hybrid_bayes_tree.dot()))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Operations on HybridBayesTree\n",
        "\n",
        "Similar to `HybridBayesNet`, the tree can be used for optimization, evaluation, and extracting specific conditional distributions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "MAP Solution (Optimize):\n",
            "HybridValues: \n",
            "  Continuous: 2 elements\n",
            "  x0: 0\n",
            "  x1: 1\n",
            "  Discrete: (d0, 0)\n",
            "  Nonlinear\n",
            "Values with 0 values:\n",
            "\n",
            "MPE Discrete Assignment:\n",
            "DiscreteValues{7205759403792793600: 0}\n"
          ]
        }
      ],
      "source": [
        "hbt = hybrid_bayes_tree # Use the tree from elimination\n",
        "\n",
        "# --- Optimization (Finding MAP state) ---\n",
        "# Computes MPE for discrete, then optimizes continuous given MPE\n",
        "map_solution = hbt.optimize()\n",
        "print(\"\\nMAP Solution (Optimize):\")\n",
        "map_solution.print()\n",
        "\n",
        "# --- MPE (Most Probable Explanation for Discrete Variables) ---\n",
        "mpe_assignment = hbt.mpe()\n",
        "print(\"\\nMPE Discrete Assignment:\")\n",
        "print(mpe_assignment)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "id": "389541c5",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Optimized Continuous Solution for D0=0:\n",
            "VectorValues: 2 elements\n",
            "  x0: 0\n",
            "  x1: 1\n",
            "\n",
            "Optimized Continuous Solution for D0=1:\n",
            "VectorValues: 2 elements\n",
            "  x0: 1.90476\n",
            "  x1: 3.09524\n",
            "\n",
            "Total Error at MAP solution: 0.023412453920796855\n"
          ]
        }
      ],
      "source": [
        "# --- Optimize Continuous given specific Discrete Assignment ---\n",
        "dv0 = DiscreteValues()\n",
        "dv0[D(0)] = 0\n",
        "cont_solution_d0_eq_0 = hbt.optimize(dv0)\n",
        "print(\"\\nOptimized Continuous Solution for D0=0:\")\n",
        "cont_solution_d0_eq_0.print()\n",
        "\n",
        "dv1 = DiscreteValues()\n",
        "dv1[D(0)] = 1\n",
        "cont_solution_d0_eq_1 = hbt.optimize(dv1)\n",
        "print(\"\\nOptimized Continuous Solution for D0=1:\")\n",
        "cont_solution_d0_eq_1.print()\n",
        "\n",
        "# --- Evaluation ---\n",
        "# Evaluate error (sum of errors of clique conditionals)\n",
        "total_error = hbt.error(map_solution)\n",
        "print(f\"\\nTotal Error at MAP solution: {total_error}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "id": "c1bd1d19",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "GaussianBayesTree for D0=0:\n",
            ": cliques: 3, variables: 2\n",
            "- p()\n",
            "  R = Empty (0x0)\n",
            "  d = Empty (0x1)\n",
            "  mean: 0 elements\n",
            "  logNormalizationConstant: -0\n",
            "  No noise model\n",
            "| - p(x1)\n",
            "  R = [ 2.21565 ]\n",
            "  d = [ 2.21565 ]\n",
            "  mean: 1 elements\n",
            "  x1: 1\n",
            "  logNormalizationConstant: -0.123394\n",
            "  No noise model\n",
            "| | - p(x0 | x1)\n",
            "  R = [ 3.31662 ]\n",
            "  S[x1] = [ -3.01511 ]\n",
            "  d = [ -3.01511 ]\n",
            "  logNormalizationConstant: 0.280009\n",
            "  No noise model\n"
          ]
        }
      ],
      "source": [
        "# --- Extract Conditional Distributions ---\n",
        "# Choose a specific GaussianBayesTree for a discrete assignment\n",
        "gaussian_bayes_tree_d0_eq_0 = hbt.choose(dv0)\n",
        "print(\"\\nGaussianBayesTree for D0=0:\")\n",
        "gaussian_bayes_tree_d0_eq_0.print()\n",
        "# display(graphviz.Source(gaussian_bayes_tree_d0_eq_0.dot()))"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "py312",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.6"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}
